import cv2
import numpy as np

def read_detect_result(file_path):
    frame_idx_2_badminton_coord = {}
    with open(file_path, 'r') as fin:
        for line in fin.readlines():
            if line.startswith('Frame'):
                print("==== ignore first frame ", line)
                continue
            frame_idx, visible, x, y = [int(_) for _ in line.strip().split(',')]
            frame_idx_2_badminton_coord[frame_idx] = [visible, x, y]
    return frame_idx_2_badminton_coord

def process_video(video_path, result_path, output_path):
    """
    使用OpenCV逐帧处理视频
    """
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print("无法打开视频文件")
        return
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 编码器
    fps = 30.0  # 帧率
    frame_size = (3840, 2160)  # 帧尺寸 (宽度, 高度)
    
    # 创建VideoWriter对象
    out = cv2.VideoWriter(output_path, fourcc, fps, frame_size)

    frame_idx_2_badminton_coord = read_detect_result(result_path)
    frame_count = 0
    
    while True:
        # 读取一帧
        ret, frame = cap.read()
        
        if frame_count % 10 == 0:
            print("finish ", frame_count)
        # 如果读取失败（视频结束）
        if not ret:
            break
        
        visible_coord = frame_idx_2_badminton_coord[frame_count]
        # 在这里进行帧处理
        processed_frame = process_frame(frame, frame_count, visible_coord)
        
        out.write(processed_frame)
        frame_count += 1
        #print(f"处理第 {frame_count} 帧")
    
    # 释放资源
    cap.release()
    out.release()

def process_frame(frame, frame_count, visible_coord):
    """
    帧处理函数 - 这里可以添加你的处理逻辑
    """
    cv2.putText(frame, str(frame_count), (100, 100), cv2.FONT_HERSHEY_SIMPLEX, 5, (0, 0, 255), 2)
    if visible_coord[0] == 1:
        coord = visible_coord[1:]
        cv2.circle(frame, coord, 10, (0, 255, 0), 2)
    return frame


if __name__ == '__main__':
    case_idx = 4
    #view = 13
    for view in [0, 4, 9, 13, 18, 22, 27, 32]:
        video_path = 'select_cases/case_%d_v_%d.mp4' % (case_idx, view)
        result_path = 'badminton_cases/case_%d_v_%d_ball.csv' % (case_idx, view)
        out_path = 'result/show_case_%d_v_%d.mp4' % (case_idx, view)
        process_video(video_path, result_path, out_path)

